
import os
import tensorflow as tf
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.layers import Dense, Conv2D, Flatten, MaxPool2D, BatchNormalization, Activation, GlobalAvgPool2D
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import SparseCategoricalCrossentropy
from tensorflow.keras.metrics import SparseCategoricalAccuracy
from tensorflow.keras.callbacks import ModelCheckpoint
import matplotlib.pyplot as plt

(train_x, train_y), (test_x, test_y) = cifar10.load_data()

train_x, test_x = train_x/255.0, test_x/255.0
print(train_x.shape, train_y.shape, test_x.shape, test_y.shape)


class ResnetBlock(Model):
    def __init__(self, filters, strides=1, residual_path=False):
        super(ResnetBlock, self).__init__()
        self.filters = filters
        self.strides = strides
        self.residual_path = residual_path

        self.c1 = Conv2D(filters, (3, 3), strides=strides, padding='same')
        self.bn1 = BatchNormalization()
        self.ac1 = Activation('relu')

        self.c2 = Conv2D(filters, (3, 3), strides=1, padding='same')
        self.bn2 = BatchNormalization()

        if residual_path:
            self.down_c1 = Conv2D(filters, (1, 1), strides=strides, padding='same', use_bias=False)
            self.down_bn = BatchNormalization()

        self.ac2 = Activation('relu')

    def call(self, inputs):
        residual = inputs

        x = self.c1(inputs)
        x = self.bn1(x)
        x = self.ac1(x)

        x = self.c2(x)
        y = self.bn2(x)

        if self.residual_path:
            residual = self.down_c1(inputs)
            residual = self.down_bn(residual)
        out = self.ac2(y + residual)
        return out


class ResnetModel(Model):
    def __init__(self):
        super(ResnetModel, self).__init__()
        self.cv1 = Conv2D(64, (3, 3), padding='same', use_bias=False)
        self.bn1 = BatchNormalization()
        self.ac1 = Activation('relu')

        self.blocks = Sequential()
        self.blocks.add(ResnetBlock(filters=64, strides=1, residual_path=False))
        self.blocks.add(ResnetBlock(filters=64, strides=1, residual_path=False))

        self.blocks.add(ResnetBlock(filters=128, strides=2, residual_path=True))
        self.blocks.add(ResnetBlock(filters=128, strides=1, residual_path=False))

        self.blocks.add(ResnetBlock(filters=256, strides=2, residual_path=True))
        self.blocks.add(ResnetBlock(filters=256, strides=1, residual_path=False))

        self.blocks.add(ResnetBlock(filters=512, strides=2, residual_path=True))
        self.blocks.add(ResnetBlock(filters=512, strides=1, residual_path=False))

        self.gap = GlobalAvgPool2D()
        self.fc = Dense(10, activation='softmax')

    def call(self, inputs, training=None, mask=None):
        x = self.cv1(inputs)
        x = self.bn1(x)
        x = self.ac1(x)

        x = self.blocks(x)
        x = self.gap(x)
        y = self.fc(x)
        return y


model = ResnetModel()
model.compile(optimizer=Adam(),
              loss=SparseCategoricalCrossentropy(from_logits=False),
              metrics=[SparseCategoricalAccuracy()])


# save model
model_save_path = './checkpoint/baseline.ckpt'
model_save_callback = ModelCheckpoint(
    model_save_path,
    monitor='val_loss',
    verbose=0,
    save_best_only=True,
    save_weights_only=True,
    mode='auto',
    save_freq='epoch'
)
# load model
if os.path.exists(model_save_path + '.index'):
    print('*************** load model *********************************')
    model.load_weights(model_save_path)


history = model.fit(train_x, train_y, batch_size=32, epochs=5, validation_data=(test_x, test_y),
                    validation_freq=1,
                    callbacks=[model_save_callback])

model.summary()

# ************************ 保存权重 *************
mode_weight_txt = './weight.txt'
with open(mode_weight_txt, 'w') as f:
    for v in model.trainable_variables:
        f.write(str(v.name) + '\n')
        f.write(str(v.shape) + '\n')
        f.write(str(v.numpy()) + '\n')

# ************ plot loass/acc curve ****************************
loss, val_loss = history.history['loss'], history.history['val_loss']
acc, val_acc = history.history['sparse_categorical_accuracy'], history.history['val_sparse_categorical_accuracy']

plt.subplot(2, 1, 1)
plt.plot(loss, label='train_loss')
plt.plot(val_loss, label='val_loss')
plt.title('loss')
plt.legend()

plt.subplot(2, 1, 2)
plt.plot(acc, label='train_acc')
plt.plot(val_acc, label='val_acc')
plt.title('accuracy')
plt.legend()

plt.show()
